'''
使用方法
1、先用VoiceToImg.py批量生成语谱图
2、执行ImgToDataset.py将语谱图打包成image.train.FRecord训练集和image.test.FRecord测试集
'''
import tensorflow as tf
import os
import random
import cv2
import numpy as np

def GetFileList():
    pathName = "按人分组\\"
    ret_dic = {}
    for file in os.listdir(pathName):
        img_path = pathName + file  #每类图片的地址
        font_label = file
        files = []
        for imgfile in os.listdir(img_path):
            if imgfile.endswith('png'):
                img_file_path = img_path + '\\' + imgfile
                files.append(img_file_path)
        ret_dic[font_label] = files
    return ret_dic

#生成字符型的属性
def _float32_feature(value):
#	return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))
	
#生成整数型的属性
def _int64_feature(value):
	return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def changeAudioToFRecord(dic):
    filename_train = 'image.train.FRecord'
    filename_test = 'image.test.FRecord'
    
    #创建一个writer来写TFRecord文件
    writer_train = tf.python_io.TFRecordWriter(filename_train)
    writer_test = tf.python_io.TFRecordWriter(filename_test)
    index = 0
    for lable, files in dic.items():
        print(str(index) + '------------------------label', str(lable))
        random.shuffle(files)
        count = len(files)
        tranCount = (int)(count * 0.7) 
        for i in range(tranCount):
            if i%10 == 0:
                print(str(i) + '/' + str(count))
            np_data = cv2.imdecode(np.fromfile(files[i],dtype=np.uint8),-1)
            np_arr = np_data.reshape(-1)
            example=tf.train.Example(features=tf.train.Features(feature={
                    'label': _int64_feature(int(lable)),
                    'aud_raw': _float32_feature(np_arr.tolist())}))
            #将一个Example写入TFRecord文件中
            writer_train.write(example.SerializeToString())
        for i in range(tranCount, count): 
            if i%10 == 0:
                print(str(i) + '/' + str(count))
            np_data = cv2.imdecode(np.fromfile(files[i],dtype=np.uint8),-1)
            np_arr = np_data.reshape(-1)
            example=tf.train.Example(features=tf.train.Features(feature={
                    'label': _int64_feature(int(lable)),
                    'aud_raw': _float32_feature(np_arr.tolist())}))
            #将一个Example写入TFRecord文件中
            writer_test.write(example.SerializeToString())
        index = index + 1
    writer_train.close()
    writer_test.close()
    return np_data

if __name__ == "__main__":
    dict = GetFileList()
    exm = changeAudioToFRecord(dict)
